#include<iostream>
#include<iomanip>
#include<bitset>
#include<vector>
#include<string>

// 定义字节类型
typedef unsigned char byte;

// S盒
byte sBox[256] = {
    // 你可以在这里定义你自己的S盒
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
};

// 加密的轮密钥
byte roundKeys[176] ;

// 轮常量
byte rcon[10] = {
    0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1B,0x36
};

// 在GF(2^8)中乘以2的运算
byte mul2(byte a) {
    // 左移一位
    byte temp = a << 1;

    // 如果最高位被移出，进行模拟2^8加上特定的多项式的操作
    if ((a & 0x80) != 0) {
        temp ^= 0x1b; // 在GF(256)上的乘法特定多项式
    }

    return temp;
}

// 在GF(2^8)中乘以3的运算
byte mul3(byte a) {
    byte temp;
    // 乘2运算
    temp = mul2(a);
    // 与原始值进行异或操作
    temp = temp ^ a;

    return temp;;
}

//密钥扩展函数
void keyExpansion(const byte* key) {
    // 前16个字节是原始密钥
    for (int i = 0; i < 16; ++i) {
        roundKeys[i] = key[i];
    }

    int bytesGenerated = 16;
    int rconIteration = 0;
    byte temp[4];

    while (bytesGenerated < 176) {
        // 从上一个轮密钥中获取4个字节
        for (int i = 0; i < 4; ++i) {
            temp[i] = roundKeys[i + bytesGenerated - 4];
        }

        // 应用密钥调度核心
        if (bytesGenerated % 16 == 0) {
            // 对4个字节执行一次循环字节替换，然后应用轮常量到第一个字节
            // 参见AES标准
            byte tempi = temp[0];
            temp[0] = sBox[temp[1]] ^ rcon[rconIteration];
            temp[1] = sBox[temp[2]];
            temp[2] = sBox[temp[3]];
            temp[3] = sBox[tempi];
            rconIteration++;
        }

        // 生成下一个轮密钥
        for (int i = 0; i < 4; ++i) {
            roundKeys[bytesGenerated] = roundKeys[bytesGenerated - 16] ^ temp[i];
            bytesGenerated++;
        }
    }
}

// 轮密钥加
void addRoundKey(byte* state, int round) {
    for (int i = 0; i < 16; ++i) {
        state[i] ^= roundKeys[round * 16 + i];
    }
}

// 字节替换
void subBytes(byte* state) {
    for (int i = 0; i < 16; ++i) {
        state[i] = sBox[state[i]];
    }
}

// 行移位
void shiftRows(byte* state) {
    byte temp[16];

    // 第一行不变
    temp[0] = state[0];
    temp[1] = state[5];
    temp[2] = state[10];
    temp[3] = state[15];

    // 第二行左移一位
    temp[4] = state[4];
    temp[5] = state[9];
    temp[6] = state[14];
    temp[7] = state[3];

    // 第三行左移两位
    temp[8] = state[8];
    temp[9] = state[13];
    temp[10] = state[2];
    temp[11] = state[7];

    // 第四行左移三位
    temp[12] = state[12];
    temp[13] = state[1];
    temp[14] = state[6];
    temp[15] = state[11];

    for (int i = 0; i < 16; ++i) {
        state[i] = temp[i];
    }
}

// 列混淆
void mixColumns(byte* state) {
    byte temp[16];

    for (int i = 0; i < 16; ++i) {
        temp[i] = state[i];
    }
    for (int i = 0; i < 4; i++) {
        state[4 * i + 0] = mul2(temp[4 * i + 0]) ^ mul3(temp[4 * i + 1]) ^ temp[4 * i + 2] ^ temp[4 * i + 3];
        state[4 * i + 1] = temp[4 * i + 0] ^ mul2(temp[4 * i + 1]) ^ mul3(temp[4 * i + 2]) ^ temp[4 * i + 3];
        state[4 * i + 2] = temp[4 * i + 0] ^ temp[4 * i + 1] ^ mul2(temp[4 * i + 2]) ^ mul3(temp[4 * i + 3]);
        state[4 * i + 3] = mul3(temp[4 * i + 0]) ^ temp[4 * i + 1] ^ temp[4 * i + 2] ^ mul2(temp[4 * i + 3]);
    }
}

// AES-128加密
void aes_encrypt(const byte* key, const byte* plaintext, byte* ciphertext) {
    keyExpansion(key);

    byte state[16];

    // 初始化状态矩阵
    for (int i = 0; i < 16; ++i) {
        state[i] = plaintext[i];
    }

    // 初始轮密钥加
    addRoundKey(state, 0);

    for (int round = 1; round < 10; ++round) {
        subBytes(state);
        shiftRows(state);
        mixColumns(state);
        addRoundKey(state, round);
    }

    subBytes(state);
    shiftRows(state);
    addRoundKey(state, 10);

    // 输出密文
    for (int i = 0; i < 16; ++i) {
        ciphertext[i] = state[i];
    }
}

std::vector<byte> hexStringToByteArray(const std::string& hexString) {
    std::vector<byte> byteArray;

    // 每两个字符转换为一个字节
    for (size_t i = 0; i < hexString.size(); i += 2) {
        std::string byteString = hexString.substr(i, 2);
        byte byteValue = static_cast<byte>(std::stoul(byteString, nullptr, 16));
        byteArray.push_back(byteValue);
    }

    return byteArray;
}

int AES(int a,int b) {
    std::string inputkey;
    std::string inputplaintext;

    byte key[16];
    byte plaintext[16];
    byte ciphertext[16];

    // 输入32位的字符串
    std::cin >> inputkey;
    std::vector<byte> tempkey = hexStringToByteArray(inputkey);
    for (int i = 0; i < 16; i++) {
        key[i] = tempkey[i];
    }
    std::cin >> inputplaintext;
    std::vector<byte> tempplaintext = hexStringToByteArray(inputplaintext);
    for (int i = 0; i < 16; i++) {
        plaintext[i] = tempplaintext[i];
    }

    aes_encrypt(key, plaintext, ciphertext);

    // 输出密文
    for (int i = 0; i < 16; ++i) {
        std::cout << std::hex << std::setw(2) << std::setfill('0') << std::uppercase << (int)ciphertext[i];
    }

    return ciphertext;
}
